﻿using System;
using System.Collections.Generic;
using System.Linq;

namespace BucaUtils
{
    public static class RecursiveJoinExtensions
    {
        public static IEnumerable<TTResult> RecursiveJoin<TSource, TKey, TTResult>(this IEnumerable<TSource> source,
           Func<TSource, TKey> parentKeySelector,
           Func<TSource, TKey> childKeySelector,
           Func<TSource, IEnumerable<TTResult>, TTResult> resultSelector)
        {
            return RecursiveJoin(source, parentKeySelector, childKeySelector,
               resultSelector, Comparer<TKey>.Default);
        }

        public static IEnumerable<TTResult> RecursiveJoin<TSource, TKey, TTResult>(this IEnumerable<TSource> source,
           Func<TSource, TKey> parentKeySelector,
           Func<TSource, TKey> childKeySelector,
           Func<TSource, int, IEnumerable<TTResult>, TTResult> resultSelector)
        {
            return RecursiveJoin(source, parentKeySelector, childKeySelector,
               (TSource element, int depth, int index, IEnumerable<TTResult> children)
                  => resultSelector(element, index, children));
        }

        public static IEnumerable<TTResult> RecursiveJoin<TSource, TKey, TTResult>(this IEnumerable<TSource> source,
           Func<TSource, TKey> parentKeySelector,
           Func<TSource, TKey> childKeySelector,
           Func<TSource, IEnumerable<TTResult>, TTResult> resultSelector,
           IComparer<TKey> comparer)
        {
            return RecursiveJoin(source, parentKeySelector, childKeySelector,
               (TSource element, int depth, int index, IEnumerable<TTResult> children)
                  => resultSelector(element, children), comparer);
        }

        public static IEnumerable<TTResult> RecursiveJoin<TSource, TKey, TTResult>(this IEnumerable<TSource> source,
           Func<TSource, TKey> parentKeySelector,
           Func<TSource, TKey> childKeySelector,
           Func<TSource, int, IEnumerable<TTResult>, TTResult> resultSelector,
           IComparer<TKey> comparer)
        {
            return RecursiveJoin(source, parentKeySelector, childKeySelector,
               (TSource element, int depth, int index, IEnumerable<TTResult> children)
                  => resultSelector(element, index, children), comparer);
        }

        public static IEnumerable<TTResult> RecursiveJoin<TSource, TKey, TTResult>(this IEnumerable<TSource> source,
           Func<TSource, TKey> parentKeySelector,
           Func<TSource, TKey> childKeySelector,
           Func<TSource, int, int, IEnumerable<TTResult>, TTResult> resultSelector)
        {
            return RecursiveJoin(source, parentKeySelector, childKeySelector,
               resultSelector, Comparer<TKey>.Default);
        }

        public static IEnumerable<TTResult> RecursiveJoin<TSource, TKey, TTResult>(this IEnumerable<TSource> source,
           Func<TSource, TKey> parentKeySelector,
           Func<TSource, TKey> childKeySelector,
           Func<TSource, int, int, IEnumerable<TTResult>, TTResult> resultSelector,
           IComparer<TKey> comparer)
        {
            // prevent source being enumerated more than once per RecursiveJoin call
            source = new LinkedList<TSource>(source);

            // fast binary search lookup
            var parents = new SortedDictionary<TKey, TSource>(comparer);
            var children
               = new SortedDictionary<TKey, LinkedList<TSource>>(comparer);

            foreach (TSource element in source)
            {
                parents[parentKeySelector(element)] = element;

                LinkedList<TSource> list;

                TKey childKey = childKeySelector(element);

                if (!children.TryGetValue(childKey, out list))
                {
                    children[childKey] = list = new LinkedList<TSource>();
                }

                list.AddLast(element);
            }

            // initialize to null otherwise compiler complains at single line assignment
            Func<TSource, int, IEnumerable<TTResult>> childSelector = null;

            childSelector = (parent, depth) =>
            {
                LinkedList<TSource> innerChildren;

                if (children.TryGetValue(parentKeySelector(parent), out innerChildren))
                {
                    return innerChildren.Select((child, index)
                       => resultSelector(child, index, depth, childSelector(child, depth + 1)));
                }

                return Enumerable.Empty<TTResult>();
            };

            var index1 = -1;
            foreach (TSource element in source)
            {
                if (!parents.ContainsKey(childKeySelector(element)))
                {
                    index1++;
                    yield return resultSelector(element, index1, 0, childSelector(element, 1));
                }
            }
        }
    }
}
